import torch.nn as nn

from monai.losses import DiceLoss, FocalLoss


class LossBraTS(nn.Module):
    def __init__(self, focal=True, redun=False):
        super(LossBraTS, self).__init__()
        self.dice = DiceLoss(sigmoid=True, batch=True)
        self.ce = FocalLoss(gamma=2.0, to_onehot_y=False) if focal else nn.BCEWithLogitsLoss()
        
        self.redundancy = redun

    def forward(self, p, y):
        return self.dice(p, y) + self.ce(p, y.float())

    def _loss(self, p, y):
        y_wt, y_tc, y_et = y > 0, ((y == 1) + (y == 3)) > 0, y == 3
        p_wt, p_tc, p_et = p[:, 0].unsqueeze(1), p[:, 1].unsqueeze(1), p[:, 2].unsqueeze(1)
        l_wt, l_tc, l_et = self._loss(p_wt, y_wt), self._loss(p_tc, y_tc), self._loss(p_et, y_et)
        
        if self.redundancy:
            return l_wt + l_tc + l_et, [l_wt, l_tc, l_et]
        else:
            return l_wt + l_tc + l_et
    
    
def TestLoss():
    # ********  Test loss func inputs and outputs
    import torch
    
    rec_x = torch.randn((2, 4, 256, 256), requires_grad=True)
    rec_y = torch.randn((2, 4, 256, 256), requires_grad=True)
    
    con_x = torch.randn((2, 4, 512), requires_grad=True)
    con_y = torch.randn((2, 4, 512), requires_grad=True)
    
    loss = Loss(1, 1, 1, 1)
    
    a, b, c, d = loss(rec_x, rec_y, con_x, con_y) 
    print(a, b, c, d)
    
    
def TestBraTSLoss():
    import numpy as np
    def trans_brats_label(x):
            mask_WT = x.copy()
            mask_WT[mask_WT == 1] = 1
            mask_WT[mask_WT == 2] = 1
            mask_WT[mask_WT == 3] = 1

            mask_TC = x.copy()
            mask_TC[mask_TC == 1] = 1
            mask_TC[mask_TC == 2] = 0
            mask_TC[mask_TC == 3] = 1

            mask_ET = x.copy()
            mask_ET[mask_ET == 1] = 0
            mask_ET[mask_ET == 2] = 0
            mask_ET[mask_ET == 3] = 1
            
            mask = np.stack([mask_WT, mask_TC, mask_ET], axis=0)
            return mask
    
    path = '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00127-000/BraTS-GLI-00127-000-seg.nii.gz'
    
    import SimpleITK as sitk
    label = sitk.ReadImage(path)
    label = sitk.GetArrayFromImage(label)
    
    brats_label = trans_brats_label(label)

    import torch
    label = torch.tensor(label)
    brats_label = torch.tensor(brats_label)
    
    pred = brats_label[:, 36].unsqueeze(0)
    # true = label[40]
    true = label[45].unsqueeze(0).unsqueeze(0)
    
    brats_loss = LossBraTS(True)
    loss = LossBraTS(focal=True)
    
    loss_1 = brats_loss(pred, true)
    loss_2 = brats_loss(pred, true)
    
    print(loss_1, loss_2)
    


if __name__ == "__main__":
    pass